"""
数据存储代理
管理多个存储实现，支持数据备份和读写分离
"""
from typing import Dict, List, Any, Optional
from fin_senti_entity_platform.data_collection.storage.base_storage import BaseStorage, StorageFactory
from fin_senti_entity_platform.utils.logger import Logger
from fin_senti_entity_platform.utils.config_loader import ConfigLoader


class StorageProxy(BaseStorage):
    """数据存储代理类"""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        初始化存储代理
        
        Args:
            config: 存储配置参数，如果为None则从配置文件加载
        """
        # 如果没有提供配置，则从配置文件加载
        if config is None:
            config = ConfigLoader().get('data_storage', {})
            
        super().__init__(config)
        self.logger = Logger().get_logger(__name__)
        self.storages = {}
        self._init_storages()
    
    def _init_storages(self):
        """初始化所有配置的存储实例"""
        storage_configs = self.config.get('storages', {})
        
        for storage_type, config in storage_configs.items():
            try:
                self.storages[storage_type] = StorageFactory.create_storage(storage_type, config)
                self.logger.info(f"成功初始化 {storage_type} 存储")
            except Exception as e:
                self.logger.error(f"初始化 {storage_type} 存储失败: {str(e)}")
    
    def get_storage(self, storage_type: str) -> Optional[BaseStorage]:
        """
        获取指定类型的存储实例
        
        Args:
            storage_type: 存储类型
        
        Returns:
            Optional[BaseStorage]: 存储实例，如果不存在则返回None
        """
        return self.storages.get(storage_type)
    
    def save(self, data: Any, collection_name: str, **kwargs) -> bool:
        """
        保存数据到主存储，并可选备份到其他存储
        
        Args:
            data: 要保存的数据
            collection_name: 集合/表名称
            **kwargs: 其他参数
                - backup_to: 要备份到的存储类型列表
        
        Returns:
            bool: 保存是否成功
        """
        # 获取主存储
        primary_storage_type = self.config.get('primary_storage', 'mongodb')
        primary_storage = self.get_storage(primary_storage_type)
        
        if not primary_storage:
            self.logger.error(f"主存储 {primary_storage_type} 未初始化")
            return False
            
        # 保存到主存储
        success = primary_storage.save(data, collection_name, **kwargs)
        
        # 如果保存成功，并且需要备份，则备份到其他存储
        if success and 'backup_to' in kwargs:
            for backup_type in kwargs['backup_to']:
                backup_storage = self.get_storage(backup_type)
                if backup_storage:
                    try:
                        backup_storage.save(data, collection_name, **kwargs)
                        self.logger.info(f"成功备份数据到 {backup_type} 存储")
                    except Exception as e:
                        self.logger.error(f"备份数据到 {backup_type} 存储失败: {str(e)}")
                        
        return success
    
    def batch_save(self, data_list: List[Any], collection_name: str, **kwargs) -> bool:
        """
        批量保存数据到主存储，并可选备份到其他存储
        
        Args:
            data_list: 数据列表
            collection_name: 集合/表名称
            **kwargs: 其他参数
                - backup_to: 要备份到的存储类型列表
        
        Returns:
            bool: 批量保存是否成功
        """
        # 获取主存储
        primary_storage_type = self.config.get('primary_storage', 'mongodb')
        primary_storage = self.get_storage(primary_storage_type)
        
        if not primary_storage:
            self.logger.error(f"主存储 {primary_storage_type} 未初始化")
            return False
            
        # 批量保存到主存储
        success = primary_storage.batch_save(data_list, collection_name, **kwargs)
        
        # 如果保存成功，并且需要备份，则备份到其他存储
        if success and 'backup_to' in kwargs:
            for backup_type in kwargs['backup_to']:
                backup_storage = self.get_storage(backup_type)
                if backup_storage:
                    try:
                        backup_storage.batch_save(data_list, collection_name, **kwargs)
                        self.logger.info(f"成功批量备份数据到 {backup_type} 存储")
                    except Exception as e:
                        self.logger.error(f"批量备份数据到 {backup_type} 存储失败: {str(e)}")
                        
        return success
    
    def query(self, collection_name: str, query_params: Dict[str, Any], **kwargs) -> List[Any]:
        """
        从主存储或指定存储查询数据
        
        Args:
            collection_name: 集合/表名称
            query_params: 查询参数
            **kwargs: 其他参数
                - from_storage: 从哪个存储查询
        
        Returns:
            List[Any]: 查询结果列表
        """
        # 确定查询的存储类型
        storage_type = kwargs.get('from_storage', self.config.get('primary_storage', 'mongodb'))
        storage = self.get_storage(storage_type)
        
        if not storage:
            self.logger.error(f"存储 {storage_type} 未初始化")
            return []
            
        # 执行查询
        return storage.query(collection_name, query_params, **kwargs)
    
    def update(self, collection_name: str, query_params: Dict[str, Any], update_data: Dict[str, Any], **kwargs) -> bool:
        """
        更新主存储中的数据
        
        Args:
            collection_name: 集合/表名称
            query_params: 查询条件
            update_data: 更新数据
            **kwargs: 其他参数
        
        Returns:
            bool: 更新是否成功
        """
        # 获取主存储
        primary_storage_type = self.config.get('primary_storage', 'mongodb')
        primary_storage = self.get_storage(primary_storage_type)
        
        if not primary_storage:
            self.logger.error(f"主存储 {primary_storage_type} 未初始化")
            return False
            
        # 执行更新
        return primary_storage.update(collection_name, query_params, update_data, **kwargs)
    
    def delete(self, collection_name: str, query_params: Dict[str, Any], **kwargs) -> bool:
        """
        删除主存储中的数据
        
        Args:
            collection_name: 集合/表名称
            query_params: 删除条件
            **kwargs: 其他参数
        
        Returns:
            bool: 删除是否成功
        """
        # 获取主存储
        primary_storage_type = self.config.get('primary_storage', 'mongodb')
        primary_storage = self.get_storage(primary_storage_type)
        
        if not primary_storage:
            self.logger.error(f"主存储 {primary_storage_type} 未初始化")
            return False
            
        # 执行删除
        return primary_storage.delete(collection_name, query_params, **kwargs)
    
    def close(self):
        """关闭所有存储连接"""
        for storage_type, storage in self.storages.items():
            try:
                storage.close()
                self.logger.info(f"已关闭 {storage_type} 存储连接")
            except Exception as e:
                self.logger.error(f"关闭 {storage_type} 存储连接失败: {str(e)}")


# 创建全局存储代理实例，方便使用
storage_proxy = StorageProxy()